package com.tianji.aigc.memory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 基于redis的会话记忆处理类
 */
@Component
public class RedisChatMemory implements ChatMemory {

//    @Autowired
//    private MessageChatMemoryAdvisor messageChatMemoryAdvisor;

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 会话完成添加消息
     *
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        //拼接redisKey
        String key = "CHAT:" + conversationId;

        //将消息转换为字符串类型
        List<String> list = messages.stream().map(message -> {
            //拷贝Message对象为自定义的RedisMessage对象
            RedisMessage redisMessage = BeanUtil.toBean(message, RedisMessage.class);
            //设置消息内容字符串，否则为空
            redisMessage.setTextContent(message.getText());
            //设置Tool相关属性的值
            if (message instanceof AssistantMessage assistantMessage) {
                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            } else if (message instanceof ToolResponseMessage toolResponseMessage) {
                redisMessage.setToolResponses(toolResponseMessage.getResponses());
            }
            return JSONUtil.toJsonStr(redisMessage);
        }).toList();


        //将消息存入redis的list中
        redisTemplate.opsForList().leftPushAll(key, list);

    }

    /**
     * 获取对话消息
     *
     * @param conversationId 对话id
     * @param lastN          消息的数量
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        //拼接redisKey
        String key = "CHAT:" + conversationId;
        //从缓存中查询消息字符串列表
        List<String> list = redisTemplate.opsForList().range(key, 0, lastN);

        //将消息字符串转换为消息对象
        List<Message> messageList = list.stream().map(m -> {
            //将消息JSON字符串转为RedisMessage对象
            RedisMessage redisMessage = JSONUtil.toBean(m, RedisMessage.class);
            Message message = null;
            //根据消息类型封装具体消息对象
            if (redisMessage.getMessageType().equals(MessageType.USER)) {
                message =  new UserMessage(redisMessage.getMessageType(),redisMessage.getTextContent(),redisMessage.getMedia(),redisMessage.getMetadata());
            } else if (redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                message =  new AssistantMessage(redisMessage.getTextContent(),redisMessage.getMetadata(),redisMessage.getToolCalls(), redisMessage.getMedia());
            } else if(redisMessage.getMessageType().equals(MessageType.SYSTEM)) {
                message = new SystemMessage(redisMessage.getTextContent());
            } else if(redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message =  new ToolResponseMessage(redisMessage.getToolResponses(),redisMessage.getMetadata());
            }
            return message;
        }).toList();
        return messageList;
    }

    /**
     * 清除对话消息
     *
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        //拼接redisKey
        String key = "CHAT:" + conversationId;
        redisTemplate.unlink(key);
    }
}
